#include <iostream>
#include <thread>
#include <vector>


// 矩阵乘法,thread库实现
// C = alpha*A * B + beta*C

typedef struct{
    int start;
    int end;
    float alpha;
    float beta;
    int k;
    float *a;
    float *b;
    float *c;
}Pot;

void mul_thread(Pot *p){
    
    int start = p->start;
    int end = p->end;
    float alpha = p->alpha;
    float beta = p->beta;
    int k = p->k;
    float *A = p->a;
    float *B = p->b;
    float *C = p->c;

    for(int i = start; i < end; i++) {
        int row = i / k;
        int col = i % k;
        C[i] = beta * C[i];
        for (int j = 0; j < k; j++) {
            C[i] += alpha * A[row * k + j] * B[j * k + col];
        }
    }

}

int dgemm(int out, float *a, int lda, float alpha, float *b, int ldb, float *c, int ldc, float beta){
    int nums = lda * out; // 矩阵元素个数
    int thread_num = std::thread::hardware_concurrency(); // 线程数量
    std::vector<std::thread> threads(thread_num);

    std::vector<Pot> pots(thread_num);

    int blocksdim;
    if (nums%thread_num != 0) {
        if (nums < thread_num) {
            for (int i = 0; i < nums; i++) {
                int end = i + 1;
                int start = i;
                pots[i] = Pot{start, end, alpha, beta, ldb, a, b, c};
            }
        }else{
            blocksdim = nums/(thread_num -1);
            for (int i = 0; i < thread_num; i++) {
                int end;
                if (i == thread_num - 1) {
                    end = nums;
                }
                else{
                    end = blocksdim * (i + 1);
                }
                int start = blocksdim * i;
                
                pots[i] = Pot{start, end, alpha, beta, ldb, a, b, c};
            }
        }
    }else{
        blocksdim = nums/thread_num;

        for (int i = 0; i < thread_num; i++) {
                int end = blocksdim * (i + 1);
                int start = blocksdim * i;
                pots[i] = Pot{start, end, alpha, beta, ldb, a, b, c};
            }
    }
    if (nums < thread_num) {
        thread_num = nums;
    }

    for (int i = 0; i < thread_num; i++) {
        threads[i] = std::thread(mul_thread, &pots[i]);
    }

    for (int i = 0; i < thread_num; i++) {
        threads[i].join();
    }

    return 0;
}

// test
int main(){
    int m = 3;
    int n = 3;
    int k = 3;

    float alpha = 1.0f;
    float beta = 0.0f;
    float A[m * k] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f};
    float B[k * n] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f};
    float C[m * n] = {0.0f};

    dgemm(n, A, m, alpha, B, k, C, m, beta);

    for(int i = 0; i < m; i++) {
        for(int j = 0; j < n; j++) {
            std::cout << C[i * n + j] << " ";
        }
        std::cout << std::endl;
    }
}